import os

import torchvision
import torch
from torch.utils.data.dataset import Dataset

from PIL import Image


class PubgPostureDataset(Dataset):
    def __init__(self):
        self.train_transforms = get_train_transforms(input_size=(28, 28))
        self.images = []
        self.labels = []

        base_path = r'C:\Programs\workspace\deep_learning\data\pubg\stand_squat_lie'
        posture_classes = sorted(os.listdir(base_path))
        for cls in posture_classes:
            cls_path = os.path.join(base_path, cls)
            for image in os.listdir(cls_path):
                self.images.append(os.path.join(cls_path, image))
                self.labels.append(posture_classes.index(cls))
                # label = torch.zeros(3)
                # label[posture_classes.index(cls)] = 1
                # self.labels.append(label)

    def __len__(self):
        return len(self.images)

    def __getitem__(self, item):
        image = Image.open(self.images[item])

        for t in self.train_transforms:
            image = t(image)

        label = self.labels[item]
        return image, label


def get_train_transforms(input_size):
    return torchvision.transforms.Compose([
        # torchvision.transforms.SSDCropping(),
        # torchvision.transforms.Resize(size=(256, 256)),
        torchvision.transforms.Resize(size=input_size),
        # torchvision.transforms.ColorJitter(),
        torchvision.transforms.ToTensor(),
        # torchvision.transforms.RandomHorizontalFlip(),
        # torchvision.transforms.Normalization(),
        torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),  # 归一化
        # torchvision.transforms.AssignGTtoDefaultBox()
    ])


train_loader = torch.utils.data.DataLoader(PubgPostureDataset(), batch_size=3, shuffle=True)

# train_loader = torch.utils.data.DataLoader(
#     torchvision.datasets.MNIST('/files/', train=True, download=True,
#                                transform=torchvision.transforms.Compose([
#                                    torchvision.transforms.ToTensor(),
#                                    torchvision.transforms.Normalize(
#                                        (0.1307,), (0.3081,))
#                                ])),
#     batch_size=64, shuffle=True)
